"""
This file provides the super class describing a null model.
All classes creating null models should inherit from this one 
"""

import random
import sys 
sys.path.append("..")
import data_handler.data_handler as data_handler
import matching.match_peptides_features as matching


class NullModel:
    """
    Class to create a null model 
    """
    
    @staticmethod    
    def create_null_model(th_peptides, obs_features, *args):
        """
        Input:
        th_peptides - list of data_handler.Peptide objects 
        obs_features - list of data_handler.feature objects
        args - list of other arguments (specific to each type of null model)
        
        Output: a function that given a data_handler.Peptide object, 
        returns its probability according to the null model
        """        
        pass



class RandRTNullModel(NullModel):
    """
    Class to create a null model by following the steps below:
    - sort the peptides according to mass, 
    - divide the peptides in n_bins, shuffle retention times within each bin
    - match the peptides to the observed MS1-features 
    - return a function that gives for a certain peptide length the fraction of 
      peptides matched to MS1-features
    """
    @staticmethod    
    def load_null_model(null_model_file):
        """
        Input: filename including a header, and then on each line the 
        first two columns represent a peptide length and the probability \
        according to the null model, tab-separated. The last 
        line is interpreter as peptides >= that peptide length
        
        Output: a null model
        """
        lines = open(null_model_file).readlines()[1:]
        prob_dict = {}
        for l in lines:
            fields = l.split("\t")
            prob_dict[int(fields[0])] = float(fields[1])
        null_model = lambda p: RandRTNullModel.null_model_function(\
                prob_dict, p.sequence)       
                
        return null_model
                
                
    @staticmethod    
    def create_null_model(th_peptides, obs_features, n_bins, \
            error_tolerance_list):
        # randomize RTs
        new_peptides = RandRTNullModel.randomize_rts(th_peptides, n_bins)
        # match the new peptides to the MS1-features 
        matching.match_peptides2feat(new_peptides, obs_features, \
                error_tolerance_list)          
        # calculate the probability dictionary 
        prob_dict = RandRTNullModel.compute_probability_dict(new_peptides)
        # get the null model as a function 
        null_model = lambda p: RandRTNullModel.null_model_function(\
                prob_dict, p.sequence)    
        return null_model, prob_dict, new_peptides


    @staticmethod            
    def null_model_function(prob_dict, sequence):
        n = len(sequence)
        if n in prob_dict:
            return prob_dict[n]
        max_len = max(prob_dict.keys())
        if n > max_len:
            return prob_dict[max_len]
        return -1.0
        
          
    @staticmethod 
    def compute_probability_dict(peptides, min_matched = 5):
        """
        Input: 
        peptides  - peptides matched to MS1-features 
        min_matched - the min number of peptides of a certain length that 
                      should be matched; the peptides longer are pooled when 
                      we encounter the first peptide length with less than 
                      min_matched peptides    
        Output: a dictionary giving for each length the fraction of peptide
                matched to MS1-features 
        """
        lengths = {}
        for p in peptides:
            n = len(p.sequence)
            if n in lengths:
                lengths[n][0] += 1
            else:
                lengths[n] = [1.0, 0.0]
            if len(p.matched_features) > 0:
                lengths[n][1] += 1

        pairs = lengths.items()    
        pairs.sort(key = lambda p: p[0])
        probs = {}
        i = 0    
        while i < len(pairs):
            n = pairs[i][0]
            len_pair = pairs[i][1]
            if len_pair[1] < min_matched: 
                s_all = sum([pair[1][0] for pair in pairs[i:]])
                s_matched = sum([pair[1][1] for pair in pairs[i:]])
                probs[n] = s_matched / s_all
                i = len(pairs)
            else:
                probs[n] = len_pair[1] / len_pair[0]
                i += 1
                                
        return probs


    @staticmethod 
    def randomize_rts_win(peptides, start, stop):
        rts = [p.pred_rt for p in peptides[start:stop]]
        random.shuffle(rts)  
        for j in range(len(rts)):
            peptides[start + j].pred_rt = rts[j]  
        

    @staticmethod         
    def randomize_rts(th_peptides, n_bins):
        """
        Input: 
        th_peptides - list of data_handler.Peptide objects 
        n_bins - number of bins to randomize the retention times         
        
        Output: list of Peptide objects with randomized RTs
        """
        new_peptides = [data_handler.Peptide(p.sequence, p.pred_rt, \
                p.th_mass, p.proteins) for p in th_peptides]
        new_peptides.sort(key = lambda p:p.th_mass)
        m = len(new_peptides) / n_bins
        i = 0
        while i < len(new_peptides) - m:
            rts = RandRTNullModel.randomize_rts_win(new_peptides, i, i + m)  
            i = i + m
        RandRTNullModel.randomize_rts_win(new_peptides, i, len(new_peptides))
    
        return new_peptides
